require 'rubygems'
require 'net/https'
require 'uri'
require 'erb'


class Analysis
  attr_accessor :mu_ip, :username, :password, :posted_uuid
  def initialize(mu_ip=nil,username='admin',password='admin')
    @mu_ip = mu_ip
    @username = username
    @password = password
  end

  def configure(params)
    template = %q{<?xml version="1.0" encoding="UTF-8"?>
<mu_config version="4.0">
    <templates>
        <analysis name="<%= analysis_name %>">
            <analyzer_mode><%= mode %></analyzer_mode>
            <attacks><%= attacks %></attacks>
            <event_actions><%= event_actions %></event_actions>
            <monitors><%= monitors %></monitors>
            <%= target %>
            <%= testbed %>
        </analysis>
    </templates>
</mu_config>
    }
    erb = ERB.new(template, 0, '%<>')

    analysis_name = "Analysis"
    mode = "Client"
    attacks = target = testbed = monitors = event_actions = ""
    params.split(",").each do |p| 
      type,template_name = p.strip.split("=")
      case type
      when 'PublishedNetworkAttack', 'ScenarioMutationAttack', 'MutationAttack','DosAttack'
        attacks = get_template(type,template_name)
      when 'CommandMonitor','ConsoleMonitor','DioMonitor','RemoteLogMonitor','SnmpMonitor','SyslogMonitor'
        monitors = get_template(type,template_name)
      when 'RemoteTrigger','SyslogNotifier'
        event_actions = get_template(type,template_name)
      when 'ServerTarget'
        target = get_template(type,template_name)
      when 'PassthroughTarget'
        target = get_template(type,template_name)
        mode = "ClientServer"
      when 'Testbed'
        testbed = get_template(type,template_name)
      when 'Name'
        analysis_name = template_name
      else
        puts "Unknown type: #{type}"
        usage
      end
    end
    ['testbed','target','attacks'].each do |element|
      if eval(element).empty?
        puts "Missing element: #{element}"
        usage
      end
    end
    configuration = erb.result(binding)
    post configuration
  end

  def usage
    puts "Configuration Elements MUST include:"
    puts "Testbed, PassthroughTarget|ServerTarget"
    puts "PublishedNetworkAttack|ScenarioMutationAttack|MutationAttack|DosAttack"
    puts "\nOPTIONS:"
    puts "CommandMonitor, ConsoleMonitor, DioMonitor, RemoteLogMonitor, SnmpMonitor, SyslogMonitor"
    puts "RemoteTrigger, SyslogNotifier\n\n"
    exit
  end
  
  def get_template(type,name)
    doc = request("/api/v3/templates/export?type=#{type}&name=#{name}")
    doc =~ /<templates>(.*)<\/templates>/m and xml = $1
    if xml.nil?
      puts "Get Template Failed: Type=#{type} Name=#{name}"
      exit
    end
    return xml.gsub!(/uuid=".*?"/,'')
  end
  
  # send request to api
  def request(doc)
    uri = URI.parse("https://#{@mu_ip}/")
    resp = "";
    http = Net::HTTP.new(uri.host, uri.port)
    http.use_ssl = true
    escaped = URI.escape(doc)
    http.start do |http|
      req = Net::HTTP::Get.new("#{escaped}", {"User-Agent" => @username})
      req.basic_auth(@username, @password)
      response = http.request(req)
      resp = response.body
    end
    if (/<.+>/).match(resp) or doc == "/api/v3/templates/types"
      return resp
    else
      puts "HTTP::Get Request Error: #{resp}"
      exit
    end
  end

  # Post the analysis/template object to the muserver.
  def post(doc)
    uri = URI.parse("https://#{@mu_ip}/")
    resp = "";
    http = Net::HTTP.new(uri.host, uri.port)
    http.use_ssl = true
    http.start do |http|
      req = Net::HTTP::Post.new("/api/v3/templates/import", {"User-Agent" => @username})
      req.basic_auth(@username, @password)
      req.body = doc.to_s
      response = http.request(req)
      resp = response.body
    end
    if (/<.+>/).match(resp)
      if (/<uuid>(.+?)<\/uuid>/).match(resp)
        @posted_uuid = $1
      end
      return resp
    else
      puts "HTTP::Post Error: #{resp}"
      exit
    end
  end

  def run(uuid=@posted_uuid)
    doc = request("/api/v3/analysis/run?uuid=#{uuid}")
    if (/uuid="(.+?)"/).match(doc)
      @run_uuid = $1
      return @run_uuid
    end
    return doc
  end
end

if $0 == __FILE__
  unless ARGV[0] 
    puts "Usage: #{$0} <mu_ip> <Element>=<template_name>, <Element>=<template_name>,.."
    puts "\nexample: ruby ./analysis.rb 10.10.5.72 Name=test run with stuff, Testbed=lab, ServerTarget=dell-eth1, MutationAttack=icmp\n\n"
    Analysis.new.usage
  end
  test = Analysis.new ARGV.shift
  puts test.configure(ARGV.join("\s"))
  puts "run uuid: " + test.run
end
